//
// Created by vivi on 10/12/2017.
//

#include "HairAcceleratedSelleSimulator.h"
#include <memory>
#include <iostream>
#include <tbb/tbb.h>

namespace XRwy {
	using namespace std;
	using namespace Eigen;

	HairAcceleratedSelleSimulatorStrandSolverContainer::HairAcceleratedSelleSimulatorStrandSolverContainer(StrandView & strand_view)  {
		strand_view_ptr = &strand_view;

		spring_ptrs.clear();

		n = strand_view.all_ids.size();

		A = vector<Band7Matrix3fGroup>(n);
		L3 = L2 = L1 = L0 = L0i = U0 = U1 = U2 = vector<Matrix3f>(n);
		y = x = b = vector<Vector3f>(n);
	}

	void HairAcceleratedSelleSimulatorStrandSolverContainer::push_spring(const HairSelleSimulatorSpring &spring) {
		spring_ptrs.push_back(&spring);
	}

	HairAcceleratedSelleSimulator::HairAcceleratedSelleSimulator(const HairReader &reader,
	                                                             const HairSelleSimulatorConfiguration &conf):
			HairSelleSimulator(reader, conf) {

		allocator<HairAcceleratedSelleSimulatorStrandSolverContainer> alloc;
		solver_containers = vector<HairAcceleratedSelleSimulatorStrandSolverContainer>();
		solver_containers.reserve(this->strands.size());

		for (size_t i = 0; i < this->strands.size(); ++i)
			solver_containers.emplace_back(this->strands[i]);

		//push each spring to the right solver container
		//since spring could not across strands
		//so we only need to check one particle
		for (const auto & spring_view : this->springs)
			solver_containers[spring_view.p1->strand_id].push_spring(spring_view);
	}

	void HairAcceleratedSelleSimulator::strand_dynamic(const Eigen::Affine3f & transformation, const float t) {
		static const Matrix3f zero = Matrix3f::Zero();
		static const Matrix3f identity = Matrix3f::Identity();
		static Matrix3f a_init;

		const auto & conf = get_conf();

		a_init = (1.0 + t * conf.air_damping) * Matrix3f::Identity();

		tbb::parallel_for(size_t(0), this->strands.size(), [&] (size_t strand_counter) {
			const auto & strand_view = strands[strand_counter];
			auto & __ = solver_containers[strand_counter];

			//We first update the position and velocity for the strand root
			const size_t root_id = strand_view.all_ids[0];
			const Vector3f new_root_pos = transformation * particles[root_id].rest_pos;
			velv[root_id] = (new_root_pos - posv[root_id]) / t;
			posv[root_id] = new_root_pos;

			for (size_t i = 0; i < strand_view.all_ids.size(); ++i)
				__.b[i] = velv[strand_view.all_ids[i]] + t * conf.acc;

			for (auto & a : __.A)
				a[3] = a_init; //we only need one to clear, others is going to assign later

			for (auto & spring_ptr : __.spring_ptrs) {
				//id1 < id2
				size_t id1 = spring_ptr->p1->local_id;
				size_t id2 = spring_ptr->p2->local_id;
				size_t diff = id2 - id1;

				Vector3f impluse = spring_ptr->impulse(t);
				Matrix3f dm = spring_ptr->matrix(t);

				__.b[id1] += impluse;
				__.b[id2] -= impluse;

				__.A[id1][3] += dm;
				__.A[id2][3] += dm;
				__.A[id1][3 + diff] = -dm;
				__.A[id2][3 - diff] = -dm;
			}

			//Clear the spring forces and direction matrixes for strand root
			//After we modify the equation of strand root
			//The solution will always be velv[root_id]
			__.A[0][3] = identity;
			__.A[0][0] = __.A[0][1] = __.A[0][2]  = __.A[0][4] = __.A[0][5] = __.A[0][6] = zero;
			__.b[0] = velv[root_id];

//#ifndef NDEBUG
//			for (size_t i = 0; i < __.n; ++i) {
//				for (size_t k = 0; k < 7; ++k)
//					cout << "A[" << i << "][" << k << "]:\n" << __.A[i][k] << endl;
//			}
//
//			for (size_t i = 0; i < __.n; ++i)
//				cout << "b[" << i << "]:\n" << __.b[i] << endl;
//#endif

			for (size_t i = 0; i < __.n; ++i) {
				//compute L3
				__.L3[i] = i >= 3 ? __.A[i][0] : zero;

				//compute L2
				__.L2[i] = i >= 2 ? __.A[i][1] : zero;
				if (i >= 3)
					__.L2[i] -= __.L3[i] * __.U0[i - 3];

				//compute L1
				__.L1[i] = i >= 1 ? __.A[i][2] : zero;
				if (i >= 2)
					__.L1[i] -= __.L2[i] * __.U0[i - 2];
				if (i >= 3)
					__.L1[i] -= __.L3[i] * __.U1[i - 3];

				//compute L0
				__.L0[i] = __.A[i][3];
				if (i >= 1)
					__.L0[i] -= __.L1[i] * __.U0[i - 1];
				if (i >= 2)
					__.L0[i] -= __.L2[i] * __.U1[i - 2];
				if (i >= 3)
					__.L0[i] -= __.L3[i] * __.U2[i - 3];

				//compute L0i
				__.L0i[i] = __.L0[i].inverse();

				//compute U2
				__.U2[i] = i + 3 < __.n ? __.L0i[i] * __.A[i][6] : zero;

				//compute U1
				__.U1[i] = i + 2 < __.n ? __.A[i][5] : zero;
				if (i >= 1)
					__.U1[i] -= __.L1[i] * __.U2[i - 1];
				__.U1[i] = __.L0i[i] * __.U1[i];

				//compute U0
				__.U0[i] = i + 1 < __.n ? __.A[i][4] : zero;
				if (i >= 1)
					__.U0[i] -= __.L1[i] * __.U1[i - 1];
				if (i >= 2)
					__.U0[i] -= __.L2[i] * __.U2[i - 2];
				__.U0[i] = __.L0i[i] * __.U0[i];

				//compute y
				__.y[i] = __.b[i];
				if (i >= 1)
					__.y[i] -= __.L1[i] * __.y[i - 1];
				if (i >= 2)
					__.y[i] -= __.L2[i] * __.y[i - 2];
				if (i >= 3)
					__.y[i] -= __.L3[i] * __.y[i - 3];
				__.y[i] = __.L0i[i] * __.y[i];
			}

			for (int i = __.n - 1; i >= 0; --i) {
				__.x[i] = __.y[i];
				if (i + 1 < __.n)
					__.x[i] -= __.U0[i] * __.x[i + 1];
				if (i + 2 < __.n)
					__.x[i] -= __.U1[i] * __.x[i + 2];
				if (i + 3 < __.n)
					__.x[i] -= __.U2[i] * __.x[i + 3];
			}

			for (size_t i = 1; i < __.n; ++i) {
				const size_t id = strand_view.all_ids[i];
				velv[id] = __.x[i];
				posv[id] += t * velv[id];
			}
		});

//		for (size_t strand_counter = 0; strand_counter < this->strands.size(); ++strand_counter) {
//
//		}
	}

	HairAcceleratedSelleSimulator::~HairAcceleratedSelleSimulator() {
	}
}